package com.zhuodewen.ai.controller.springai;

import com.zhuodewen.ai.base.JSONResult;
import com.zhuodewen.ai.constant.CommonConstants;
import com.zhuodewen.ai.dto.springai.SpringAiTool;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiImageModel;
import org.springframework.ai.openai.OpenAiImageOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

/**
 * Controller : 控制器类
 * 用于对请求的内容、响应的内容进行数据格式转换。
 */
@RestController                         //等于@ResponseBody(返回JSON格式的数据) + @Controller(定义为Controller接口类)
@RequestMapping(value = "springai")     //路径映射
@Slf4j                                  //日志
public class SpringAiController {

    @Autowired
    OpenAiChatModel      openAiChatModel;
    @Autowired
    OllamaChatModel      ollamaChatModel;
    @Autowired
    EmbeddingModel       openAiEmbeddingModel;
    @Autowired
    OllamaEmbeddingModel ollamaEmbeddingModel;
    // @Autowired
    // VectorStore          vectorStore;
    @Autowired
    OpenAiImageModel     openAiImageModel;

    /**
     * deepseek对话
     *
     * @param message
     * @return
     */
    @PostMapping("deepseek/chat")
    public JSONResult deepSeekChat(@RequestParam(value = "message") String message) {
        /*ChatResponse response = chatModel.call(
                new Prompt(
                        "Generate the names of 5 famous pirates.",
                        OpenAiChatOptions.builder()
                                .model("deepseek-chat")
                                .temperature(0.4)
                                .build()
                ));*/
        String callback = openAiChatModel.call(message);
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, callback);
    }

    /**
     * deepseek对话-流式响应
     *
     * @param message
     * @return
     */
    @PostMapping(value = "deepseek/streamingChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> deepSeekStreamingChat(@RequestParam(value = "message") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return ChatClient.create(openAiChatModel).prompt(prompt).stream().content();
    }

    /**
     * ollama对话(本地模型)
     *
     * @param message
     * @return
     */
    @PostMapping("ollama/chat")
    public JSONResult ollamaChat(@RequestParam(value = "message") String message) {
        /*ChatResponse response = ollamaChatModel.call(
                new Prompt(
                        "Generate the names of 5 famous pirates.",
                        OllamaOptions.builder()
                                .model(OllamaModel.LLAMA3_1)
                                .temperature(0.4)
                                .build()
                ));*/
        String callback = ollamaChatModel.call(message);
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, callback);
    }

    /**
     * deepseek对话-流式响应
     *
     * @param message
     * @return
     */
    @PostMapping(value = "ollama/streamingChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> ollamaStreamingChat(@RequestParam(value = "message") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        return ChatClient.create(ollamaChatModel).prompt(prompt).stream().content();
    }

    /**
     * openai嵌入
     *
     * @param message
     * @return
     */
    @PostMapping("openai/embedding")
    public JSONResult openaiEmbedding(@RequestParam(value = "message") String message) {
         /*EmbeddingResponse embeddingResponse = openAiEmbeddingModel.call(
                new OpenAiEmbeddingModel(
                    this.openAiApi,
                    MetadataMode.EMBED,
                    OpenAiEmbeddingOptions.builder()
                            .model("text-embedding-ada-002")
                            .user("user-6")
                            .build(),
                    RetryUtils.DEFAULT_RETRY_TEMPLATE)
                );*/
        EmbeddingResponse embeddingResponse = openAiEmbeddingModel.embedForResponse(List.of(message));
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, Map.of("embedding", embeddingResponse));
    }

    /**
     * ollama嵌入
     *
     * @param message
     * @return
     */
    @PostMapping("ollama/embedding")
    public JSONResult ollamaEmbedding(@RequestParam(value = "message") String message) {
        /*EmbeddingResponse embeddingResponse = ollamaEmbeddingModel.call(
                new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
                        OllamaOptions.builder()
                                .model("chroma/all-minilm-l6-v2-f32"))
                        .truncate(false)
                        .build());*/
        EmbeddingResponse embeddingResponse = ollamaEmbeddingModel.embedForResponse(List.of(message));
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, Map.of("embedding", embeddingResponse));
    }

    /**
     * 加载知识库-使用Redis Stack方式
     *
     * @param message
     * @return
     */
    @PostMapping("ragTest")
    public JSONResult ragTest(@RequestParam(value = "message") String message) {
        // 1.加载文档
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2"))
        );

        // 2.将知识库存入向量数据库
        // vectorStore.add(documents);
        //
        // // 3.检索与查询相似的文档
        // List<Document> results = vectorStore.similaritySearch(
        //         SearchRequest.builder()
        //                 .query(message)
        //                 .topK(5)
        //                 .build()
        // );

        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, documents);
    }

    /**
     * 工具调用
     *
     * @param message
     * @return
     */
    @PostMapping("toolTest")
    public JSONResult toolTest(@RequestParam(value = "message") String message) {
        String content = ChatClient.create(openAiChatModel)
                .prompt()
                .user(message)
                .tools(new SpringAiTool())
                .call()
                .content();
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, content);
    }

    /**
     * 多模态-图片生成(暂不支持国内常见的多模态)
     *
     * @param message
     * @return
     */
    @PostMapping("imageTest")
    public JSONResult imageTest(@RequestParam(value = "message") String message) {
        ImageResponse response = openAiImageModel.call(
                new ImagePrompt("A light cream colored mini golden doodle",
                        OpenAiImageOptions.builder()
                                .quality("hd")
                                .N(4)
                                .height(1024)
                                .width(1024).build())

        );
        log.info("图片生成结果:{}", response);
        return new JSONResult().markSuccess(CommonConstants.RESULT_SUCCESS_MSG, response.getResult().getOutput().getUrl());
    }


}
